#encoding=utf-8
import xml.etree.ElementTree as ET
import cv2
import os
import sys
import numpy as np
import glob





def _get_bbox(annotation_file):
    class_names = ('BACKGROUND',
                   '1111',
                   )
    class_dict = {class_name: i for i, class_name in enumerate(class_names)}
    et = ET.parse(annotation_file)
    objects = et.findall("object")
    size = et.findall("size")
    w = int(size[0].find('width').text)
    h = int(size[0].find('height').text)
    bboxs = []
    labels = []
    for object in objects:
        class_name = object.find('name').text.lower().strip()
        if class_name == '1111':
            bndbox = object.find('bndbox')
            # VOC dataset format follows Matlab, in which indexes start from 0

            xmin = int(bndbox.find('xmin').text)
            ymin = int(bndbox.find('ymin').text)
            xmax = int(bndbox.find('xmax').text)
            ymax = int(bndbox.find('ymax').text)

            bboxs.append([xmin, ymin, xmax, ymax])

            labels.append(class_dict[class_name])
    return bboxs, labels




def _get_landmarks(annotation_file):
    class_names = ('BACKGROUND',
                   '1',
                   )
    class_dict = {class_name: i for i, class_name in enumerate(class_names)}
    et = ET.parse(annotation_file)
    objects = et.findall("object")
    size = et.findall("size")
    w = int(size[0].find('width').text)
    h = int(size[0].find('height').text)
    polyes = []
    labels = []
    for object in objects:
        class_name = object.find('name').text.lower().strip()
        if class_name == '1':
            polygon = object.find('polygon')
            # VOC dataset format follows Matlab, in which indexes start from 0

            point0 = polygon.find('point0').text.split(',')
            point1 = polygon.find('point1').text.split(',')
            point2 = polygon.find('point2').text.split(',')
            point3 = polygon.find('point3').text.split(',')
            point4 = polygon.find('point4').text.split(',')
            fpoint0 = [0,0]
            fpoint1 = [0,0]
            fpoint2 = [0,0]
            fpoint3 = [0,0]
            fpoint4 = [0,0]
            fpoint0[0] = int(point0[0])
            fpoint0[1] = int(point0[1])
            fpoint1[0] = int(point4[0])
            fpoint1[1] = int(point4[1])
            fpoint2[0] = int(point1[0])
            fpoint2[1] = int(point1[1])
            fpoint3[0] = int(point2[0])
            fpoint3[1] = int(point2[1])
            fpoint4[0] = int(point3[0])
            fpoint4[1] = int(point3[1])
            polyes.append([fpoint0, fpoint1, fpoint2, fpoint3, fpoint4])

            labels.append(class_dict[class_name])
    return polyes,labels

if __name__ == "__main__":
    concat_file = '/media/handsome/data/hanson/WIDER_train/others/heils-backlight-face/bboxs_landmarks_backlight.txt'
    root_path = r'/media/handsome/data/hanson/WIDER_train/others/heils-backlight-face'
    bbox_src_pth = os.path.join(root_path, "Annotations-bbox")
    lds_src_pth = os.path.join(root_path, "Annotations-landmarks")
    img_src_pth = os.path.join(root_path, "Images")
    save_src_pth = os.path.join(root_path, "Total_Images")
    f = open(concat_file, 'w')

    for idx, folder in enumerate(os.listdir(bbox_src_pth)):
        bbox_folder_pth = os.path.join(bbox_src_pth,folder)
        lds_folder_pth = os.path.join(lds_src_pth, folder)
        # print(bbox_folder_pth, lds_folder_pth)
        for jdx, bbox_xml_path in enumerate(glob.glob(os.path.join(bbox_folder_pth, '*.xml'))):
            lds_xml_path = os.path.join(lds_folder_pth,bbox_xml_path.split('/')[-1])
            if not os.path.exists(lds_xml_path):
                print(lds_xml_path)
                continue

            polyes, _ = _get_landmarks(lds_xml_path)
            bboxs, _ = _get_bbox(bbox_xml_path)
            img_pth = os.path.join(img_src_pth,lds_xml_path.split('/')[-2], lds_xml_path.split('/')[-1].replace('.xml','.jpg'))
            img = cv2.imread(img_pth)
            file_name = img_pth.split('/')[-1]

            if len(polyes) == 0 or len(bboxs) == 0:
                continue
            points = polyes[0]

            points = sorted(points, key=lambda x: x[1])
            eye = sorted([points[0],points[1]],key = lambda x: x[0])
            nose = points[2]
            mouth = sorted([points[3],points[4]], key=lambda x: x[0])

            sort_pts = [eye[0],eye[1], nose, mouth[0],mouth[1]]


            box = bboxs[0]

            f.write('%s %s %s %s %s %s %s %s %s %s %s %s %s %s %s \n' %
                    (file_name, str(box[0]), str(box[1]), str(box[2]), str(box[3]),
                     sort_pts[0][0], sort_pts[0][1], sort_pts[1][0], sort_pts[1][1], sort_pts[2][0], sort_pts[2][1],
                     sort_pts[3][0], sort_pts[3][1], sort_pts[4][0], sort_pts[4][1]))
            cv2.imwrite(os.path.join(save_src_pth,file_name),img)

    f.close()
            # cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (255, 0, 0), 2)
            # for xxi, point in enumerate(sort_pts):
            #     cv2.putText(img, str(xxi), (point[0], point[1]), cv2.FONT_HERSHEY_COMPLEX, 0.8, (0, 255, 0), 1)
            #     cv2.circle(img, (point[0], point[1]), 2, (255, 0, 0), 2)
            #
            #
            # cv2.imshow('img',img)
            # cv2.waitKey(0)


# bbox_file = '/media/handsome/data/hanson/CelebA/Anno/list_bbox_celeba.txt'
# landmarks_file = '/media/handsome/data/hanson/CelebA/Anno/list_landmarks_celeba.txt'
# concat_file = '/media/handsome/data/hanson/CelebA/Anno/bboxs_landmarks_celeba.txt'
#
# f = open(concat_file, 'w')
#
# with open(bbox_file, 'r') as f1:
#     bboxs = f1.readlines()
# with open(landmarks_file, 'r') as f2:
#     landmarks = f2.readlines()
#
#     for idx, _ in enumerate(landmarks):
#
#         if idx < 2:
#             continue
#         landmark = landmarks[idx].strip().split(' ')
#         landmark = [x for x in landmark if x != '']
#         bbox = bboxs[idx].strip().split(' ')
#         bbox = [x for x in bbox if x != '']
#         if len(bbox) != 5 or len(landmark) != 11:
#             print("error",bbox[0], landmark[0])
#         if(landmark[0] == bbox[0]):
#             f.write('%s %s %s %s %s %s %s %s %s %s %s %s %s %s %s \n'%
#                     (bbox[0], bbox[1], bbox[2], str(int(bbox[3]) + int(bbox[1])),str(int(bbox[2]) + int(bbox[4])),
#                         landmark[1], landmark[2], landmark[3], landmark[4], landmark[5], landmark[6],
#                      landmark[7], landmark[8], landmark[9], landmark[10]))
# f.close()